﻿#include "PNGReadWrite.h"

#include <png.h>

#include <string.h>

#include "Output/Console.h"
#include "File/FileUtils.h"

bool WH::ReadPNG(const SPath &InFilePath, TVector<uint8_t> &OutData, ETextureFormat &OutFormat, uint64_t &OutWidth, uint64_t &OutHeight, uint64_t &OutPitch)
{
    auto File=SFileUtils::ReadToBinary(InFilePath);
    return ReadPNGFromFileData(File, OutData, OutFormat, OutWidth, OutHeight, OutPitch);
}

bool WH::ReadPNGFromFileData(const TVector<uint8_t> & InPNGFileData, TVector<uint8_t> &OutData, ETextureFormat &OutFormat, uint64_t &OutWidth, uint64_t &OutHeight, uint64_t &OutPitch)
{
    uint8_t *FileData=InPNGFileData.Data();

    //Read header
    uint8_t Header[8];
    //fread(Header, 1, 8, File);
    memcpy(Header, FileData, 8);
    FileData += 8;

    //Check if it is png
    if (png_sig_cmp(Header, 0, 8))
    {
        WH::PrintLine(U"这不是png文件,无法读取PNG文件");
        return false;
    }

    //Create png struct
    png_structp PngStruct = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (!PngStruct)
    {
        WH::PrintLine(U"png struct 创建失败，导致读取PNG文件失败");
        return false;
    }

    //Create png info struct

    png_infop PngInfo = png_create_info_struct(PngStruct);
    if (!PngInfo)
    {
        WH::PrintLine(U"png info 创建失败，导致读取PNG文件失败");
        return false;
    }

    //Create png end info struct
    png_infop PngEndInfo = png_create_info_struct(PngStruct);
    if (!PngEndInfo)
    {
        WH::PrintLine(U"png end info 创建失败，导致读取PNG文件失败");
        return false;
    }

    //Init png reading
    //png_init_io(PngStruct, File);
    png_set_read_fn(PngStruct, &FileData, [](png_structp png_ptr, png_bytep data, png_size_t length)
    {
        memcpy(data, *(png_bytepp)png_get_io_ptr(png_ptr), length);
        *(png_bytepp)png_get_io_ptr(png_ptr) += length;

    });

    //Let libpng know you already read the first 8 bytes
    png_set_sig_bytes(PngStruct, 8);

    //Read all the info up to the image data
    png_read_info(PngStruct, PngInfo);


    //Variables to pass to get info
    int BitDepth, ColorType;
    png_uint_32 TmpWidth, TmpHeight;

    //Get info about png
    png_get_IHDR(PngStruct, PngInfo, &TmpWidth, &TmpHeight, &BitDepth, &ColorType, NULL, NULL, NULL);

    OutWidth = TmpWidth;
    OutHeight = TmpHeight;


    //Update the png info struct.
    png_read_update_info(PngStruct, PngInfo);

    //Row size in bytes.
    auto Pitch = png_get_rowbytes(PngStruct, PngInfo);

    TVector<uint8_t> TempData;
    TempData.Resize(Pitch * OutHeight);

    //Row_pointers is for pointing to image_data for reading the png with libpng
    png_bytep* RowPointers = (png_bytep*)malloc(sizeof(png_bytep) * OutHeight);
    for (int i = 0; i < OutHeight; i++)
    {
        RowPointers[i] = TempData.Data() + i * Pitch;
    }
    
    //Read the png into image_data through row_pointers
    png_read_image(PngStruct, RowPointers);

    if (ColorType == PNG_COLOR_TYPE_RGB_ALPHA)
    {
        OutFormat = ETextureFormat::R8G8B8A8;
        OutPitch = OutWidth * 4;
        OutData.Resize(OutPitch * OutHeight);

        if(BitDepth == 16)
        {

            //Convert to RGBA
            for (int i = 0; i < OutHeight; i++)
            {
                for (int j = 0; j < OutWidth; j++)
                {
                    OutData[i * OutPitch + j * 4 + 0] = TempData[i * Pitch + j * 8 + 0] + (uint16_t)(TempData[i * Pitch + j * 8 + 1] <<8) ;
                    OutData[i * OutPitch + j * 4 + 1] = TempData[i * Pitch + j * 8 + 2] + (uint16_t)(TempData[i * Pitch + j * 8 + 3] <<8) ;
                    OutData[i * OutPitch + j * 4 + 2] = TempData[i * Pitch + j * 8 + 4] + (uint16_t)(TempData[i * Pitch + j * 8 + 5] <<8) ;
                    OutData[i * OutPitch + j * 4 + 3] = TempData[i * Pitch + j * 8 + 6] + (uint16_t)(TempData[i * Pitch + j * 8 + 7] <<8) ;

                }
            }
        }
        else if(BitDepth == 8)
        {
            //Convert to RGBA
            for (int i = 0; i < OutHeight; i++)
            {
                for (int j = 0; j < OutWidth; j++)
                {
                    OutData[i * OutPitch + j * 4 + 0] = TempData[i * Pitch + j * 4 + 0];
                    OutData[i * OutPitch + j * 4 + 1] = TempData[i * Pitch + j * 4 + 1];
                    OutData[i * OutPitch + j * 4 + 2] = TempData[i * Pitch + j * 4 + 2];
                    OutData[i * OutPitch + j * 4 + 3] = TempData[i * Pitch + j * 4 + 3];
                }
            }
        }
        else
        {
            WH::PrintLine(U"尚不支持{}位的PNG图片", BitDepth);
            return false;
        }


    }
    else if(ColorType==PNG_COLOR_TYPE_RGB)
    {
        OutFormat = ETextureFormat::R8G8B8A8;
        OutPitch = OutWidth * 4;
        OutData.Resize(OutPitch * OutHeight);

        if(BitDepth == 16)
        {

            //Convert to RGBA
            for (int i = 0; i < OutHeight; i++)
            {
                for (int j = 0; j < OutWidth; j++)
                {
                    OutData[i * OutPitch + j * 4 + 0] = TempData[i * Pitch + j * 6 + 0] + (uint16_t)(TempData[i * Pitch + j * 6 + 1] <<8) ;
                    OutData[i * OutPitch + j * 4 + 1] = TempData[i * Pitch + j * 6 + 2] + (uint16_t)(TempData[i * Pitch + j * 6 + 3] <<8) ;
                    OutData[i * OutPitch + j * 4 + 2] = TempData[i * Pitch + j * 6 + 4] + (uint16_t)(TempData[i * Pitch + j * 6 + 5] <<8) ;
                    OutData[i * OutPitch + j * 4 + 3] = 255;

                }
            }
        }
        else if(BitDepth == 8)
        {
            //Convert to RGBA
            for (int i = 0; i < OutHeight; i++)
            {
                for (int j = 0; j < OutWidth; j++)
                {
                    OutData[i * OutPitch + j * 4 + 0] = TempData[i * Pitch + j * 3 + 0];
                    OutData[i * OutPitch + j * 4 + 1] = TempData[i * Pitch + j * 3 + 1];
                    OutData[i * OutPitch + j * 4 + 2] = TempData[i * Pitch + j * 3 + 2];
                    OutData[i * OutPitch + j * 4 + 3] = 255;
                }
            }
        }
        else
        {
            WH::PrintLine(U"尚不支持{}位的PNG图片", BitDepth);
            return false;
        }

    }
    else
    {
        WH::PrintLine(U"尚不支持该PNG的通道数");
        return false;
    }



    //Clean up
    free(RowPointers);
    png_destroy_read_struct(&PngStruct, &PngInfo, &PngEndInfo);

    return true;

}

bool WH::WritePNG(const SPath &InFilePath, const TVector<uint8_t> &InData, ETextureFormat InFormat, uint64_t InWidth, uint64_t InHeight, uint64_t InPitch)
{
    assert(InFormat==ETextureFormat::R8G8B8A8 && "only support R8G8B8A8 format");

    auto Path8 = InFilePath.ToString().ToUtf8();
    FILE* File = fopen((const char *)Path8.Data(), "wb");
    if (!File)
    {
        WH::PrintLine(U"无法打开文件: {}", InFilePath.ToString());
        return false;
    }

    png_structp PngStruct = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (!PngStruct)
    {
        WH::PrintLine(U"创建png struct失败，导致写入PNG文件失败: {}", InFilePath.ToString());
        return false;
    }

    png_infop PngInfo = png_create_info_struct(PngStruct);
    if (!PngInfo)
    {
        WH::PrintLine(U"创建png info struct失败，导致写入PNG文件失败: {}", InFilePath.ToString());
        return false;
    }

    png_init_io(PngStruct, File);

    //Write header
    png_set_IHDR(PngStruct, PngInfo, InWidth, InHeight, 8, PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    //Write info
    png_write_info(PngStruct, PngInfo);

    //Write image data
    png_bytep* RowPointers = (png_bytep*)malloc(sizeof(png_bytep) * InHeight);
    for (int i = 0; i < InHeight; i++)
    {
        RowPointers[i] = (png_bytep)InData.Data() + i * InPitch;
    }

    png_write_image(PngStruct, RowPointers);

    //End write
    png_write_end(PngStruct, NULL);

    //Clean up

    free(RowPointers);
    png_destroy_write_struct(&PngStruct, &PngInfo);

    fclose(File);

    return true;
}

bool WH::WritePNGToFileData(const TVector<uint8_t> &InTextureRawData, ETextureFormat InFormat, uint64_t InWidth, uint64_t InHeight, uint64_t InPitch, TVector<uint8_t> &OutPNGFileData)
{
    assert(InFormat==ETextureFormat::R8G8B8A8 && "only support R8G8B8A8 format");

    png_structp PngStruct = png_create_write_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
    if (!PngStruct)
    {
        WH::PrintLine(U"创建png struct失败，导致写入PNG文件失败");
        return false;
    }

    png_infop PngInfo = png_create_info_struct(PngStruct);
    if (!PngInfo)
    {
        WH::PrintLine(U"创建png info struct失败，导致写入PNG文件失败");
        return false;
    }

    auto PngWriteCallback= [](png_structp PngStruct, png_bytep Data, png_size_t Length)
    {
        TVector<uint8_t> *FileData = (TVector<uint8_t> *)png_get_io_ptr(PngStruct);
        
        FileData->Reserve(FileData->Num() + Length);
        for(int i=0;i<Length;i++)
        {
            FileData->Add(Data[i]);
        }
        //FileData->Append(Data, Length);
    };

    png_set_write_fn(PngStruct, &OutPNGFileData, PngWriteCallback, NULL);

    //Write header
    png_set_IHDR(PngStruct, PngInfo, InWidth, InHeight, 8, PNG_COLOR_TYPE_RGB_ALPHA, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_BASE, PNG_FILTER_TYPE_BASE);

    //Write info
    png_write_info(PngStruct, PngInfo);

    //Write image data
    png_bytep* RowPointers = (png_bytep*)malloc(sizeof(png_bytep) * InHeight);

    for (int i = 0; i < InHeight; i++)
    {
        RowPointers[i] = (png_bytep)InTextureRawData.Data() + i * InPitch;
    }

    png_write_image(PngStruct, RowPointers);

    //End write
    png_write_end(PngStruct, NULL);

    //Clean up
    free(RowPointers);
    png_destroy_write_struct(&PngStruct, &PngInfo);

    return true;
}
